// Author: 8891689
// https://github.com/8891689
// gcc -Wall -maes -O3 -march=native aes_maes.c aes_test.c -o aes_test
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "aes_maes.h" 

// --- 十六進制轉換輔助函數 ---
int hex_char_to_int(char c) {
    if (c >= '0' && c <= '9') return c - '0';
    if (c >= 'a' && c <= 'f') return c - 'a' + 10;
    if (c >= 'A' && c <= 'F') return c - 'A' + 10;
    return -1;
}

int hex_str_to_bytes(const char *hex, uint8_t *bytes, size_t max_len) {
    size_t len = strlen(hex);
    if (len % 2 != 0 || (len / 2) > max_len) {
        fprintf(stderr, "錯誤：無效的十六進制字符串長度 (%zu) 或超出最大長度 (%zu)。\n", len, max_len);
        return -1;
    }
    for (size_t i = 0; i < len; i += 2) {
        int high = hex_char_to_int(hex[i]);
        int low = hex_char_to_int(hex[i + 1]);
        if (high == -1 || low == -1) {
            fprintf(stderr, "錯誤：十六進制字符串中包含無效字符：%c%c\n", hex[i], hex[i+1]);
            return -1;
        }
        bytes[i / 2] = (uint8_t)((high << 4) | low);
    }
    return (int)(len / 2);
}

void bytes_to_hex_str(const uint8_t *bytes, size_t len, char *hex_str) {
    for (size_t i = 0; i < len; i++) {
        sprintf(hex_str + (i * 2), "%02x", bytes[i]);
    }
    hex_str[len * 2] = '\0';
}

// --- 工作模式實現  ---

// CBC 模式實現
int AES_CBC_Encrypt_Mode(uint8_t *output_ciphertext,
                         const uint8_t *input_plaintext,
                         size_t data_len_bytes,
                         const uint8_t *initialization_vector,
                         const __m128i *round_keys_encrypt,
                         int num_rounds) {
    if (data_len_bytes == 0) return 0;
    if (data_len_bytes % AES_BLOCK_SIZE != 0) {
        fprintf(stderr, "AES_CBC_Encrypt_Mode: data_len_bytes 必須是 AES_BLOCK_SIZE (%d) 的倍數。\n", AES_BLOCK_SIZE);
        return -1;
    }
    if (output_ciphertext == NULL || input_plaintext == NULL || initialization_vector == NULL || round_keys_encrypt == NULL) {
        fprintf(stderr, "AES_CBC_Encrypt_Mode: 傳入空指針。\n");
        return -2;
    }

    size_t num_blocks = data_len_bytes / AES_BLOCK_SIZE;
    uint8_t current_block_input[AES_BLOCK_SIZE];
    uint8_t previous_cipher_block[AES_BLOCK_SIZE];

    memcpy(previous_cipher_block, initialization_vector, AES_BLOCK_SIZE);

    for (size_t i = 0; i < num_blocks; ++i) {
        const uint8_t *plaintext_ptr = input_plaintext + (i * AES_BLOCK_SIZE);
        uint8_t *ciphertext_ptr = output_ciphertext + (i * AES_BLOCK_SIZE);

        for (int j = 0; j < AES_BLOCK_SIZE; ++j) {
            current_block_input[j] = plaintext_ptr[j] ^ previous_cipher_block[j];
        }
        AES_Encrypt_Block(ciphertext_ptr, current_block_input, round_keys_encrypt, num_rounds); // 調用核心庫函數
        memcpy(previous_cipher_block, ciphertext_ptr, AES_BLOCK_SIZE);
    }
    return 0;
}

int AES_CBC_Decrypt_Mode(uint8_t *output_plaintext,
                         const uint8_t *input_ciphertext,
                         size_t data_len_bytes,
                         const uint8_t *initialization_vector,
                         const __m128i *round_keys_decrypt,
                         int num_rounds) {
    if (data_len_bytes == 0) return 0;
    if (data_len_bytes % AES_BLOCK_SIZE != 0) {
        fprintf(stderr, "AES_CBC_Decrypt_Mode: data_len_bytes 必須是 AES_BLOCK_SIZE (%d) 的倍數。\n", AES_BLOCK_SIZE);
        return -1;
    }
    if (output_plaintext == NULL || input_ciphertext == NULL || initialization_vector == NULL || round_keys_decrypt == NULL) {
        fprintf(stderr, "AES_CBC_Decrypt_Mode: 傳入空指針。\n");
        return -2;
    }

    size_t num_blocks = data_len_bytes / AES_BLOCK_SIZE;
    uint8_t decrypted_temp_block[AES_BLOCK_SIZE];
    uint8_t previous_cipher_block[AES_BLOCK_SIZE];

    memcpy(previous_cipher_block, initialization_vector, AES_BLOCK_SIZE);

    for (size_t i = 0; i < num_blocks; ++i) {
        const uint8_t *ciphertext_ptr = input_ciphertext + (i * AES_BLOCK_SIZE);
        uint8_t *plaintext_ptr = output_plaintext + (i * AES_BLOCK_SIZE);

        AES_Decrypt_Block(decrypted_temp_block, ciphertext_ptr, round_keys_decrypt, num_rounds); // 調用核心庫函數
        for (int j = 0; j < AES_BLOCK_SIZE; ++j) {
            plaintext_ptr[j] = decrypted_temp_block[j] ^ previous_cipher_block[j];
        }
        memcpy(previous_cipher_block, ciphertext_ptr, AES_BLOCK_SIZE);
    }
    return 0;
}

// CTR 模式輔助函數：增加計數器
static void increment_counter(uint8_t *counter, int block_size) {
    for (int i = block_size - 1; i >= 0; --i) {
        if (++counter[i] != 0) {
            break;
        }
    }
}

// CTR 模式實現 (加密和解密過程相同)
int AES_CTR_Mode(uint8_t *output_data,
                 const uint8_t *input_data,
                 size_t data_len_bytes,
                 const uint8_t *nonce_counter, // 初始 nonce + counter 值
                 const __m128i *round_keys_encrypt,
                 int num_rounds) {
    if (data_len_bytes == 0) return 0;
    // CTR 模式可以處理非塊對齊數據，但這裡簡化為塊對齊，或需要處理最後一個不完整塊
    if (output_data == NULL || input_data == NULL || nonce_counter == NULL || round_keys_encrypt == NULL) {
         fprintf(stderr, "AES_CTR_Mode: 傳入空指針。\n");
        return -2;
    }

    uint8_t current_counter_block[AES_BLOCK_SIZE];
    uint8_t keystream_block[AES_BLOCK_SIZE];
    memcpy(current_counter_block, nonce_counter, AES_BLOCK_SIZE);

    size_t processed_bytes = 0;
    while(processed_bytes < data_len_bytes) {
        AES_Encrypt_Block(keystream_block, current_counter_block, round_keys_encrypt, num_rounds); // 調用核心庫函數

        size_t remaining_bytes = data_len_bytes - processed_bytes;
        size_t bytes_to_xor = (remaining_bytes < AES_BLOCK_SIZE) ? remaining_bytes : AES_BLOCK_SIZE;

        for (size_t j = 0; j < bytes_to_xor; ++j) {
            output_data[processed_bytes + j] = input_data[processed_bytes + j] ^ keystream_block[j];
        }

        increment_counter(current_counter_block, AES_BLOCK_SIZE);
        processed_bytes += bytes_to_xor;
    }
    return 0;
}

// --- CFB128 Mode Implementation ---
int AES_CFB128_Encrypt_Mode(uint8_t *output_ciphertext,
                             const uint8_t *input_plaintext,
                             size_t data_len_bytes,
                             const uint8_t *initialization_vector,
                             const __m128i *round_keys_encrypt,
                             int num_rounds) {
    if (data_len_bytes == 0) return 0;
    if (data_len_bytes % AES_BLOCK_SIZE != 0) {
        fprintf(stderr, "AES_CFB128_Encrypt_Mode: data_len_bytes 必須是 AES_BLOCK_SIZE (%d) 的倍數。\n", AES_BLOCK_SIZE);
        return -1;
    }
    if (output_ciphertext == NULL || input_plaintext == NULL || initialization_vector == NULL || round_keys_encrypt == NULL) {
        fprintf(stderr, "AES_CFB128_Encrypt_Mode: 傳入空指針。\n");
        return -2;
    }

    size_t num_blocks = data_len_bytes / AES_BLOCK_SIZE;
    uint8_t current_input_to_aes[AES_BLOCK_SIZE]; // 用於送入 AES 加密的塊 (IV 或前一個密文塊)
    uint8_t keystream_block[AES_BLOCK_SIZE];      // AES 加密 IV/前一個密文塊的輸出

    memcpy(current_input_to_aes, initialization_vector, AES_BLOCK_SIZE);

    for (size_t i = 0; i < num_blocks; ++i) {
        const uint8_t *plaintext_ptr = input_plaintext + (i * AES_BLOCK_SIZE);
        uint8_t *ciphertext_ptr = output_ciphertext + (i * AES_BLOCK_SIZE);

        // 1. 加密 IV 或前一個密文塊，得到キーストリーム
        AES_Encrypt_Block(keystream_block, current_input_to_aes, round_keys_encrypt, num_rounds);

        // 2. キーストリーム與明文異或得到密文
        for (int j = 0; j < AES_BLOCK_SIZE; ++j) {
            ciphertext_ptr[j] = plaintext_ptr[j] ^ keystream_block[j];
        }

        // 3. 將當前生成的密文塊作爲下一次 AES 加密的輸入
        memcpy(current_input_to_aes, ciphertext_ptr, AES_BLOCK_SIZE);
    }
    return 0;
}

int AES_CFB128_Decrypt_Mode(uint8_t *output_plaintext,
                             const uint8_t *input_ciphertext,
                             size_t data_len_bytes,
                             const uint8_t *initialization_vector,
                             const __m128i *round_keys_encrypt, // 注意：CFB解密仍使用加密密鑰
                             int num_rounds) {
    if (data_len_bytes == 0) return 0;
    if (data_len_bytes % AES_BLOCK_SIZE != 0) {
        fprintf(stderr, "AES_CFB128_Decrypt_Mode: data_len_bytes 必須是 AES_BLOCK_SIZE (%d) 的倍數。\n", AES_BLOCK_SIZE);
        return -1;
    }
    if (output_plaintext == NULL || input_ciphertext == NULL || initialization_vector == NULL || round_keys_encrypt == NULL) {
        fprintf(stderr, "AES_CFB128_Decrypt_Mode: 傳入空指針。\n");
        return -2;
    }

    size_t num_blocks = data_len_bytes / AES_BLOCK_SIZE;
    uint8_t current_input_to_aes[AES_BLOCK_SIZE]; // 用於送入 AES 加密的塊 (IV 或前一個密文塊)
    uint8_t keystream_block[AES_BLOCK_SIZE];      // AES 加密 IV/前一個密文塊的輸出

    memcpy(current_input_to_aes, initialization_vector, AES_BLOCK_SIZE);

    for (size_t i = 0; i < num_blocks; ++i) {
        const uint8_t *ciphertext_ptr = input_ciphertext + (i * AES_BLOCK_SIZE);
        uint8_t *plaintext_ptr = output_plaintext + (i * AES_BLOCK_SIZE);

        // 1. 加密 IV 或前一個密文塊，得到キーストリーム
        AES_Encrypt_Block(keystream_block, current_input_to_aes, round_keys_encrypt, num_rounds);

        // 2. キーストリーム與密文異或得到明文
        for (int j = 0; j < AES_BLOCK_SIZE; ++j) {
            plaintext_ptr[j] = ciphertext_ptr[j] ^ keystream_block[j];
        }

        // 3. 將當前輸入的密文塊作爲下一次 AES 加密的輸入
        memcpy(current_input_to_aes, ciphertext_ptr, AES_BLOCK_SIZE);
    }
    return 0;
}


// --- OFB Mode Implementation ---
int AES_OFB_Mode(uint8_t *output_data,
                 const uint8_t *input_data,
                 size_t data_len_bytes,
                 const uint8_t *initialization_vector,
                 const __m128i *round_keys_encrypt,
                 int num_rounds) {
    if (data_len_bytes == 0) return 0;
    // OFB 模式可以處理非塊對齊數據，但這裡簡化為塊對齊，或需要處理最後一個不完整塊
    if (output_data == NULL || input_data == NULL || initialization_vector == NULL || round_keys_encrypt == NULL) {
         fprintf(stderr, "AES_OFB_Mode: 傳入空指針。\n");
        return -2;
    }

    uint8_t current_input_to_aes[AES_BLOCK_SIZE]; // 用於送入 AES 加密的塊 (IV 或前一個AES輸出塊)
    uint8_t keystream_block[AES_BLOCK_SIZE];      // AES 加密 IV/前一個AES輸出塊的結果

    memcpy(current_input_to_aes, initialization_vector, AES_BLOCK_SIZE);

    size_t processed_bytes = 0;
    while(processed_bytes < data_len_bytes) {
        // 1. 加密 IV 或前一個 AES 的輸出，得到キーストリーム塊
        AES_Encrypt_Block(keystream_block, current_input_to_aes, round_keys_encrypt, num_rounds);

        size_t remaining_bytes = data_len_bytes - processed_bytes;
        size_t bytes_to_xor = (remaining_bytes < AES_BLOCK_SIZE) ? remaining_bytes : AES_BLOCK_SIZE;

        // 2. キーストリーム與輸入數據異或
        for (size_t j = 0; j < bytes_to_xor; ++j) {
            output_data[processed_bytes + j] = input_data[processed_bytes + j] ^ keystream_block[j];
        }

        // 3. 將當前生成的キーストリーム塊 (即AES的輸出) 作爲下一次 AES 加密的輸入
        memcpy(current_input_to_aes, keystream_block, AES_BLOCK_SIZE);
        processed_bytes += bytes_to_xor;
    }
    return 0;
}
// --- 測試函數定義 ---
void test_single_block_aes() {
    printf("\n--- 單塊 AES 操作測試 ---\n");
    AES_KeySize key_type_enum = AES_KEY_128;
    int num_rounds = AES_Nr_Rounds(key_type_enum);
    uint8_t key_bytes[16] = {
        0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
        0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c
    };
    uint8_t plaintext_block[AES_BLOCK_SIZE] = {
        0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d,
        0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37, 0x07, 0x34
    };
    uint8_t expected_ciphertext_block[AES_BLOCK_SIZE] = {
        0x39, 0x25, 0x84, 0x1d, 0x02, 0xdc, 0x09, 0xfb,
        0xdc, 0x11, 0x85, 0x97, 0x19, 0x6a, 0x0b, 0x32
    };

    __m128i round_keys_encrypt[15]; // 足夠 AES-256
    __m128i round_keys_decrypt[15];
    uint8_t ciphertext_block[AES_BLOCK_SIZE];
    uint8_t decrypted_block[AES_BLOCK_SIZE];

    printf("明文: "); AES_PrintState(plaintext_block);

    if (AES_KeyExpansion(key_bytes, round_keys_encrypt, key_type_enum, AES_KEY_SCHEDULE_ENCRYPT) != 0) {
        fprintf(stderr, "單塊加密密鑰擴展失敗\n");
        return;
    }
    AES_Encrypt_Block(ciphertext_block, plaintext_block, round_keys_encrypt, num_rounds);
    printf("密文 (實際): "); AES_PrintState(ciphertext_block);
    printf("密文 (期望): "); AES_PrintState(expected_ciphertext_block);

    if (memcmp(ciphertext_block, expected_ciphertext_block, AES_BLOCK_SIZE) == 0) {
        printf("單塊加密: 成功!\n");
    } else {
        printf("單塊加密: 失敗!\n");
    }

    if (AES_KeyExpansion(key_bytes, round_keys_decrypt, key_type_enum, AES_KEY_SCHEDULE_DECRYPT) != 0) {
        fprintf(stderr, "單塊解密密鑰擴展失敗\n");
        return;
    }
    AES_Decrypt_Block(decrypted_block, ciphertext_block, round_keys_decrypt, num_rounds);
    printf("解密後: "); AES_PrintState(decrypted_block);

    if (memcmp(decrypted_block, plaintext_block, AES_BLOCK_SIZE) == 0) {
        printf("單塊解密: 成功!\n");
    } else {
        printf("單塊解密: 失敗!\n");
    }
    printf("---------------------------\n");
}

void benchmark_avx2_aes() {
    // 這個函數測試的是 ECB 模式 (通過 AES_Encrypt_AVX2) 的性能
    AES_KeySize key_type_enum = AES_KEY_256;
    int key_len_bytes = key_type_enum * 4;
    int num_rounds = AES_Nr_Rounds(key_type_enum);

    const size_t data_size = 16 * 1024 * 1024;
    const int iterations = 100;
    size_t num_blocks = data_size / AES_BLOCK_SIZE;

    uint8_t *key_bytes = NULL;
    uint8_t *plaintext = NULL;
    uint8_t *ciphertext_avx2 = NULL;
    uint8_t *decryptedtext_avx2 = NULL;
    __m128i *w_encrypt_avx2 = NULL;
    __m128i *w_decrypt_avx2 = NULL;

    // --- 將這些變量聲明提前 ---
    clock_t start_encrypt_avx2 = 0; // 初始化爲0
    clock_t end_encrypt_avx2 = 0;
    clock_t start_decrypt_avx2 = 0;
    clock_t end_decrypt_avx2 = 0;

    double total_data_mb_avx2 = 0.0; // 初始化爲0.0
    double encrypt_duration_avx2 = 0.0;
    double encrypt_speed_mbps_avx2 = 0.0;
    double decrypt_duration_avx2 = 0.0;
    double decrypt_speed_mbps_avx2 = 0.0;
    // --- 聲明提前結束 ---


    key_bytes = (uint8_t*)malloc(key_len_bytes);
    if (!key_bytes) {fprintf(stderr, "malloc key_bytes failed\n"); return;} // 簡單返回，或者將清理邏輯放在這裡
    plaintext = (uint8_t*)malloc(data_size);
    if (!plaintext) {fprintf(stderr, "malloc plaintext failed\n"); free(key_bytes); return;}
    ciphertext_avx2 = (uint8_t*)malloc(data_size);
    if (!ciphertext_avx2) {fprintf(stderr, "malloc ciphertext_avx2 failed\n"); free(key_bytes); free(plaintext); return;}
    decryptedtext_avx2 = (uint8_t*)malloc(data_size);
    if (!decryptedtext_avx2) {fprintf(stderr, "malloc decryptedtext_avx2 failed\n"); free(key_bytes); free(plaintext); free(ciphertext_avx2); return;}


    w_encrypt_avx2 = (__m128i*)_mm_malloc((num_rounds + 1) * sizeof(__m128i), 32);
    if (!w_encrypt_avx2) {fprintf(stderr, "_mm_malloc w_encrypt_avx2 failed\n"); goto cleanup_avx2_full;}
    w_decrypt_avx2 = (__m128i*)_mm_malloc((num_rounds + 1) * sizeof(__m128i), 32);
    if (!w_decrypt_avx2) {fprintf(stderr, "_mm_malloc w_decrypt_avx2 failed\n"); goto cleanup_avx2_full;}


    srand(time(NULL));
    for (int i = 0; i < key_len_bytes; ++i) key_bytes[i] = rand() % 256;
    for (size_t i = 0; i < data_size; ++i) plaintext[i] = rand() % 256;

    if (AES_KeyExpansion(key_bytes, w_encrypt_avx2, key_type_enum, AES_KEY_SCHEDULE_ENCRYPT) != 0) {
        fprintf(stderr, "AVX2 基準測試加密密鑰擴展失敗\n");
        goto cleanup_avx2_full;
    }
    if (AES_KeyExpansion(key_bytes, w_decrypt_avx2, key_type_enum, AES_KEY_SCHEDULE_DECRYPT) != 0) {
        fprintf(stderr, "AVX2 基準測試解密密鑰擴展失敗\n");
        goto cleanup_avx2_full;
    }

    printf("\n開始 AVX2 加密基準測試 (AES-%d ECB 模式)...\n", key_type_enum * 32);
    start_encrypt_avx2 = clock(); // 現在是賦值，不是聲明+初始化
    for (int i = 0; i < iterations; ++i) {
        memcpy(ciphertext_avx2, plaintext, data_size);
        AES_Encrypt_AVX2(ciphertext_avx2, num_blocks, w_encrypt_avx2, num_rounds);
    }
    end_encrypt_avx2 = clock(); // 賦值
    printf("AVX2 加密完成。\n");

    printf("開始 AVX2 解密基準測試 (AES-%d ECB 模式)...\n", key_type_enum * 32);
    start_decrypt_avx2 = clock(); // 賦值
    for (int i = 0; i < iterations; ++i) {
        memcpy(decryptedtext_avx2, ciphertext_avx2, data_size);
        AES_Decrypt_AVX2(decryptedtext_avx2, num_blocks, w_decrypt_avx2, num_rounds);
    }
    end_decrypt_avx2 = clock(); // 賦值
    printf("AVX2 解密完成。\n");

    if (memcmp(plaintext, decryptedtext_avx2, data_size) != 0) {
        printf("\nAVX2 基準測試驗證: 失敗!\n");
        int error_count = 0;
        for (size_t offset = 0; offset < data_size && offset < 5 * AES_BLOCK_SIZE; offset += AES_BLOCK_SIZE) {
             if (memcmp(plaintext + offset, decryptedtext_avx2 + offset, AES_BLOCK_SIZE) != 0) {
                error_count++;
                if (error_count <= 5) {
                    printf("在數據塊 %zu 處不匹配\n", offset / AES_BLOCK_SIZE);
                    printf("  明文: "); AES_PrintState(plaintext + offset);
                    printf("  解密後: "); AES_PrintState(decryptedtext_avx2 + offset);
                }
            }
        }
         if (error_count == 0 && data_size > 5 * AES_BLOCK_SIZE) { // 如果前面都對，但總體錯了
             printf("前 %d 個數據塊匹配，但後續數據塊存在不匹配。\n", 5);
        } else if (error_count > 0 && error_count <=5 && num_blocks > (size_t)error_count) { // 如果打印了部分錯誤
             printf("...以及 %zu 個數據塊中可能還有更多不匹配。\n", num_blocks);
        } else if (error_count == 0 && num_blocks <= 5) { // 如果數據很少且都對，但總體驗證失敗
             printf("所有檢查的數據塊均匹配，但總體驗證失敗，可能存在未打印的錯誤。\n");
        }
    } else {
        printf("\nAVX2 基準測試驗證: 成功!\n");
    }

    // --- 計算結果 ---
    total_data_mb_avx2 = (double)data_size * iterations / (1024 * 1024);
    encrypt_duration_avx2 = (double)(end_encrypt_avx2 - start_encrypt_avx2) / CLOCKS_PER_SEC;
    encrypt_speed_mbps_avx2 = (encrypt_duration_avx2 > 0) ? (total_data_mb_avx2 / encrypt_duration_avx2) : 0;
    decrypt_duration_avx2 = (double)(end_decrypt_avx2 - start_decrypt_avx2) / CLOCKS_PER_SEC;
    decrypt_speed_mbps_avx2 = (decrypt_duration_avx2 > 0) ? (total_data_mb_avx2 / decrypt_duration_avx2) : 0;

    printf("\nAVX2 AES Standard measurement test results (ECB model)\n");
    printf("--------------------------------------------------------\n");
    printf("mode                   : AES-%d-ECB (AVX2/AES-NI)\n", key_type_enum * 32);
    printf("Total data volume      : %.2f MB (Common %d next transfer)\n", total_data_mb_avx2, iterations);
    printf("--------------------------------------------------------\n");
    printf("Encryption  time       : %.4f 秒\n", encrypt_duration_avx2);
    printf("Encryption speed       : %.2f MB/秒\n", encrypt_speed_mbps_avx2);
    printf("\n");
    printf("Decryption  time       : %.4f 秒\n", decrypt_duration_avx2);
    printf("Decryption speed       : %.2f MB/秒\n", decrypt_speed_mbps_avx2);
    printf("--------------------------------------------------------\n");

cleanup_avx2_full: // 標籤用於完整的清理
    free(key_bytes); // key_bytes 總是被分配，所以總是可以 free (如果是 NULL 也安全)
    free(plaintext);
    free(ciphertext_avx2);
    free(decryptedtext_avx2);
    if (w_encrypt_avx2) _mm_free(w_encrypt_avx2);
    if (w_decrypt_avx2) _mm_free(w_decrypt_avx2);
}


void test_cbc_mode() {
    printf("\n--- AES CBC 模式測試 ---\n");
    AES_KeySize key_type_enum = AES_KEY_128;
    int num_rounds = AES_Nr_Rounds(key_type_enum);

    uint8_t key_bytes[16] = {
        0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
        0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c
    };
    uint8_t iv[AES_BLOCK_SIZE] = {
        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
        0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f
    };
    uint8_t plaintext_cbc[2 * AES_BLOCK_SIZE] = {
        0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96,
        0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93, 0x17, 0x2a,
        0xae, 0x2d, 0x8a, 0x57, 0x1e, 0x03, 0xac, 0x9c,
        0x9e, 0xb7, 0x6f, 0xac, 0x45, 0xaf, 0x8e, 0x51
    };
    uint8_t expected_ciphertext_cbc[2 * AES_BLOCK_SIZE] = {
        0x76, 0x49, 0xab, 0xac, 0x81, 0x19, 0xb2, 0x46,
        0xce, 0xe9, 0x8e, 0x9b, 0x12, 0xe9, 0x19, 0x7d,
        0x50, 0x86, 0xcb, 0x9b, 0x50, 0x72, 0x19, 0xee,
        0x95, 0xdb, 0x11, 0x3a, 0x91, 0x76, 0x78, 0xb2
    };

    size_t data_len = 2 * AES_BLOCK_SIZE;
    uint8_t ciphertext_cbc[2 * AES_BLOCK_SIZE];
    uint8_t decrypted_cbc[2 * AES_BLOCK_SIZE];

    __m128i round_keys_encrypt[15];
    __m128i round_keys_decrypt[15];

    printf("CBC 模式明文 (2個數據塊):\n");
    AES_PrintState(plaintext_cbc);
    AES_PrintState(plaintext_cbc + AES_BLOCK_SIZE);

    if (AES_KeyExpansion(key_bytes, round_keys_encrypt, key_type_enum, AES_KEY_SCHEDULE_ENCRYPT) != 0) {
        fprintf(stderr, "CBC 加密密鑰擴展失敗\n"); return;
    }
    if (AES_CBC_Encrypt_Mode(ciphertext_cbc, plaintext_cbc, data_len, iv, round_keys_encrypt, num_rounds) != 0) {
        fprintf(stderr, "AES_CBC_Encrypt_Mode 執行失敗\n"); return;
    }

    printf("CBC 模式密文 (實際):\n"); AES_PrintState(ciphertext_cbc); AES_PrintState(ciphertext_cbc + AES_BLOCK_SIZE);
    printf("CBC 模式密文 (期望):\n"); AES_PrintState(expected_ciphertext_cbc); AES_PrintState(expected_ciphertext_cbc + AES_BLOCK_SIZE);

    if (memcmp(ciphertext_cbc, expected_ciphertext_cbc, data_len) == 0) printf("CBC 加密: 成功!\n");
    else printf("CBC 加密: 失敗!\n");

    if (AES_KeyExpansion(key_bytes, round_keys_decrypt, key_type_enum, AES_KEY_SCHEDULE_DECRYPT) != 0) {
        fprintf(stderr, "CBC 解密密鑰擴展失敗\n"); return;
    }
    if (AES_CBC_Decrypt_Mode(decrypted_cbc, ciphertext_cbc, data_len, iv, round_keys_decrypt, num_rounds) != 0) {
        fprintf(stderr, "AES_CBC_Decrypt_Mode 執行失敗\n"); return;
    }
    printf("CBC 模式解密後:\n"); AES_PrintState(decrypted_cbc); AES_PrintState(decrypted_cbc + AES_BLOCK_SIZE);

    if (memcmp(decrypted_cbc, plaintext_cbc, data_len) == 0) printf("CBC 解密: 成功!\n");
    else printf("CBC 解密: 失敗!\n");
    printf("---------------------------\n");
}

void test_ctr_mode() {
    printf("\n--- AES CTR 模式測試 (NIST SP 800-38A F.5.1) ---\n");
    AES_KeySize key_type_enum = AES_KEY_128;
    int num_rounds = AES_Nr_Rounds(key_type_enum);

    uint8_t key_bytes[16] = { // F.5.1 Key
        0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
        0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c
    };
    // F.5.1 Init. Counter (This is T1, the first counter block to be encrypted)
    uint8_t nonce_counter_initial[AES_BLOCK_SIZE] = {
        0xf0, 0xf1, 0xf2, 0xf3, 0xf4, 0xf5, 0xf6, 0xf7,
        0xf8, 0xf9, 0xfa, 0xfb, 0xfc, 0xfd, 0xfe, 0xff
    };
    // 兩個數據塊的明文 (F.5.1 Plaintext)
    uint8_t plaintext_ctr[2 * AES_BLOCK_SIZE] = {
        // Block 1
        0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96,
        0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93, 0x17, 0x2a,
        // Block 2
        0xae, 0x2d, 0x8a, 0x57, 0x1e, 0x03, 0xac, 0x9c,
        0x9e, 0xb7, 0x6f, 0xac, 0x45, 0xaf, 0x8e, 0x51
    };
    // 期望的密文 (F.5.1 Ciphertext)
    uint8_t expected_ciphertext_ctr[2 * AES_BLOCK_SIZE] = {
        // Ciphertext Block #1
        0x87, 0x4d, 0x61, 0x91, 0xb6, 0x20, 0xe3, 0x26,
        0x1b, 0xef, 0x68, 0x64, 0x99, 0x0d, 0xb6, 0xce,
        // Ciphertext Block #2
        0x98, 0x06, 0xf6, 0x6b, 0x79, 0x70, 0xfd, 0xff,
        0x86, 0x17, 0x18, 0x7b, 0xb9, 0xff, 0xfd, 0xff
    };

    size_t data_len = 2 * AES_BLOCK_SIZE;
    uint8_t ciphertext_ctr[2 * AES_BLOCK_SIZE];
    uint8_t decrypted_ctr[2 * AES_BLOCK_SIZE];

    __m128i round_keys_encrypt[15]; // CTR mode only uses encryption

    printf("CTR 模式明文 (2個數據塊):\n");
    AES_PrintState(plaintext_ctr);
    AES_PrintState(plaintext_ctr + AES_BLOCK_SIZE);
    printf("CTR 模式初始計數器 T1:\n");
    AES_PrintState(nonce_counter_initial);


    if (AES_KeyExpansion(key_bytes, round_keys_encrypt, key_type_enum, AES_KEY_SCHEDULE_ENCRYPT) != 0) {
        fprintf(stderr, "CTR 加密密鑰擴展失敗\n"); return;
    }

    // CTR 加密
    // 為了完全匹配NIST的T1, T2序列，我們需要確保AES_CTR_Mode使用的計數器與NIST一致
    // NIST的T2是f0f1f2f3f4f5f6f7f8f9fafbfcfdff00，是T1的最後一個字節ff變成00，前面字節進位
    uint8_t current_nonce_counter_enc[AES_BLOCK_SIZE];
    memcpy(current_nonce_counter_enc, nonce_counter_initial, AES_BLOCK_SIZE);
    if (AES_CTR_Mode(ciphertext_ctr, plaintext_ctr, data_len, current_nonce_counter_enc, round_keys_encrypt, num_rounds) != 0) {
        fprintf(stderr, "AES_CTR_Mode 加密執行失敗\n"); return;
    }

    printf("CTR 模式密文 (實際):\n"); AES_PrintState(ciphertext_ctr); AES_PrintState(ciphertext_ctr + AES_BLOCK_SIZE);
    printf("CTR 模式密文 (期望):\n"); AES_PrintState(expected_ciphertext_ctr); AES_PrintState(expected_ciphertext_ctr + AES_BLOCK_SIZE);

    if (memcmp(ciphertext_ctr, expected_ciphertext_ctr, data_len) == 0) printf("CTR 加密: 成功!\n");
    else printf("CTR 加密: 失敗!\n");

    // CTR 解密 (使用相同的函數和初始計數器)
    uint8_t current_nonce_counter_dec[AES_BLOCK_SIZE];
    memcpy(current_nonce_counter_dec, nonce_counter_initial, AES_BLOCK_SIZE); // 重新獲取初始計數器
    if (AES_CTR_Mode(decrypted_ctr, ciphertext_ctr, data_len, current_nonce_counter_dec, round_keys_encrypt, num_rounds) != 0) {
        fprintf(stderr, "AES_CTR_Mode 解密執行失敗\n"); return;
    }
    printf("CTR 模式解密後:\n"); AES_PrintState(decrypted_ctr); AES_PrintState(decrypted_ctr + AES_BLOCK_SIZE);

    if (memcmp(decrypted_ctr, plaintext_ctr, data_len) == 0) printf("CTR 解密: 成功!\n");
    else printf("CTR 解密: 失敗!\n");
    printf("---------------------------\n");
}

void test_cfb128_mode() {
    printf("\n--- AES CFB128 模式測試 (NIST SP 800-38A F.3.13 & F.3.14) ---\n");
    AES_KeySize key_type_enum = AES_KEY_128;
    int num_rounds = AES_Nr_Rounds(key_type_enum);

    uint8_t key_bytes[16] = { // F.3.13 Key
        0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
        0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c
    };
    uint8_t iv[AES_BLOCK_SIZE] = { // F.3.13 IV
        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
        0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f
    };
    // 兩個數據塊的明文 (來自NIST通用明文，但CFB示例只用了2塊)
    uint8_t plaintext_cfb[2 * AES_BLOCK_SIZE] = {
        // Block 1
        0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96,
        0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93, 0x17, 0x2a,
        // Block 2
        0xae, 0x2d, 0x8a, 0x57, 0x1e, 0x03, 0xac, 0x9c,
        0x9e, 0xb7, 0x6f, 0xac, 0x45, 0xaf, 0x8e, 0x51
    };
    // 期望的密文 (F.3.13 Ciphertext Segments 1 & 2)
    uint8_t expected_ciphertext_cfb[2 * AES_BLOCK_SIZE] = {
        // Ciphertext Segment #1
        0x3b, 0x3f, 0xd9, 0x2e, 0xb7, 0x2d, 0xad, 0x20,
        0x33, 0x34, 0x49, 0xf8, 0xe8, 0x3c, 0xfb, 0x4a,
        // Ciphertext Segment #2
        0xc8, 0xa6, 0x45, 0x37, 0xa0, 0xb3, 0xa9, 0x3f,
        0xcd, 0xe3, 0xcd, 0xad, 0x9f, 0x1c, 0xe5, 0x8b
    };

    size_t data_len = 2 * AES_BLOCK_SIZE;
    uint8_t ciphertext_cfb[2 * AES_BLOCK_SIZE];
    uint8_t decrypted_cfb[2 * AES_BLOCK_SIZE];

    __m128i round_keys_encrypt[15]; // CFB模式加解密都使用加密密鑰

    printf("CFB128 模式明文 (2個數據塊):\n");
    AES_PrintState(plaintext_cfb);
    AES_PrintState(plaintext_cfb + AES_BLOCK_SIZE);

    if (AES_KeyExpansion(key_bytes, round_keys_encrypt, key_type_enum, AES_KEY_SCHEDULE_ENCRYPT) != 0) {
        fprintf(stderr, "CFB128 加密密鑰擴展失敗\n"); return;
    }

    if (AES_CFB128_Encrypt_Mode(ciphertext_cfb, plaintext_cfb, data_len, iv, round_keys_encrypt, num_rounds) != 0) {
        fprintf(stderr, "AES_CFB128_Encrypt_Mode 執行失敗\n"); return;
    }

    printf("CFB128 模式密文 (實際):\n"); AES_PrintState(ciphertext_cfb); AES_PrintState(ciphertext_cfb + AES_BLOCK_SIZE);
    printf("CFB128 模式密文 (期望):\n"); AES_PrintState(expected_ciphertext_cfb); AES_PrintState(expected_ciphertext_cfb + AES_BLOCK_SIZE);

    if (memcmp(ciphertext_cfb, expected_ciphertext_cfb, data_len) == 0) printf("CFB128 加密: 成功!\n");
    else printf("CFB128 加密: 失敗!\n");

    // CFB128 解密
    // 注意：CFB解密也使用加密方向的密鑰擴展
    if (AES_CFB128_Decrypt_Mode(decrypted_cfb, ciphertext_cfb, data_len, iv, round_keys_encrypt, num_rounds) != 0) {
        fprintf(stderr, "AES_CFB128_Decrypt_Mode 執行失敗\n"); return;
    }
    printf("CFB128 模式解密後:\n"); AES_PrintState(decrypted_cfb); AES_PrintState(decrypted_cfb + AES_BLOCK_SIZE);

    if (memcmp(decrypted_cfb, plaintext_cfb, data_len) == 0) printf("CFB128 解密: 成功!\n");
    else printf("CFB128 解密: 失敗!\n");
    printf("---------------------------\n");
}

void test_ofb_mode() {
    printf("\n--- AES OFB 模式測試 (NIST SP 800-38A F.4.1 & F.4.2) ---\n");
    AES_KeySize key_type_enum = AES_KEY_128;
    int num_rounds = AES_Nr_Rounds(key_type_enum);

    uint8_t key_bytes[16] = { // F.4.1 Key
        0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
        0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c
    };
    uint8_t iv[AES_BLOCK_SIZE] = { // F.4.1 IV
        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
        0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f
    };
    // 兩個數據塊的明文 (來自NIST通用明文，但OFB示例只用了2塊)
    uint8_t plaintext_ofb[2 * AES_BLOCK_SIZE] = {
        // Block 1
        0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96,
        0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93, 0x17, 0x2a,
        // Block 2
        0xae, 0x2d, 0x8a, 0x57, 0x1e, 0x03, 0xac, 0x9c,
        0x9e, 0xb7, 0x6f, 0xac, 0x45, 0xaf, 0x8e, 0x51
    };
    // 期望的密文 (F.4.1 Ciphertext Blocks 1 & 2)
    uint8_t expected_ciphertext_ofb[2 * AES_BLOCK_SIZE] = {
        // Ciphertext Block #1
        0x3b, 0x3f, 0xd9, 0x2e, 0xb7, 0x2d, 0xad, 0x20,
        0x33, 0x34, 0x49, 0xf8, 0xe8, 0x3c, 0xfb, 0x4a,
        // Ciphertext Block #2
        0x77, 0x89, 0x50, 0x8d, 0x16, 0x91, 0x8f, 0x03,
        0xf5, 0x3c, 0x52, 0xda, 0xc5, 0x4e, 0xd8, 0x25
    };

    size_t data_len = 2 * AES_BLOCK_SIZE;
    uint8_t ciphertext_ofb[2 * AES_BLOCK_SIZE];
    uint8_t decrypted_ofb[2 * AES_BLOCK_SIZE];

    __m128i round_keys_encrypt[15]; // OFB模式加解密都使用加密密鑰

    printf("OFB 模式明文 (2個數據塊):\n");
    AES_PrintState(plaintext_ofb);
    AES_PrintState(plaintext_ofb + AES_BLOCK_SIZE);

    if (AES_KeyExpansion(key_bytes, round_keys_encrypt, key_type_enum, AES_KEY_SCHEDULE_ENCRYPT) != 0) {
        fprintf(stderr, "OFB 加密密鑰擴展失敗\n"); return;
    }

    // OFB 加密
    uint8_t iv_enc[AES_BLOCK_SIZE];
    memcpy(iv_enc, iv, AES_BLOCK_SIZE); // OFB 會修改 IV 副本
    if (AES_OFB_Mode(ciphertext_ofb, plaintext_ofb, data_len, iv_enc, round_keys_encrypt, num_rounds) != 0) {
        fprintf(stderr, "AES_OFB_Mode 加密執行失敗\n"); return;
    }

    printf("OFB 模式密文 (實際):\n"); AES_PrintState(ciphertext_ofb); AES_PrintState(ciphertext_ofb + AES_BLOCK_SIZE);
    printf("OFB 模式密文 (期望):\n"); AES_PrintState(expected_ciphertext_ofb); AES_PrintState(expected_ciphertext_ofb + AES_BLOCK_SIZE);

    if (memcmp(ciphertext_ofb, expected_ciphertext_ofb, data_len) == 0) printf("OFB 加密: 成功!\n");
    else printf("OFB 加密: 失敗!\n");

    // OFB 解密 (使用相同的函數和初始IV)
    uint8_t iv_dec[AES_BLOCK_SIZE];
    memcpy(iv_dec, iv, AES_BLOCK_SIZE); // 重新獲取初始IV
    if (AES_OFB_Mode(decrypted_ofb, ciphertext_ofb, data_len, iv_dec, round_keys_encrypt, num_rounds) != 0) {
        fprintf(stderr, "AES_OFB_Mode 解密執行失敗\n"); return;
    }
    printf("OFB 模式解密後:\n"); AES_PrintState(decrypted_ofb); AES_PrintState(decrypted_ofb + AES_BLOCK_SIZE);

    if (memcmp(decrypted_ofb, plaintext_ofb, data_len) == 0) printf("OFB 解密: 成功!\n");
    else printf("OFB 解密: 失敗!\n");
    printf("---------------------------\n");
}
int main() {
    test_single_block_aes();
    test_cbc_mode();
    test_ctr_mode();
    test_cfb128_mode(); 
    test_ofb_mode();    
    benchmark_avx2_aes(); // 函數測試的是 ECB 模式的性能 (通過 AES_Encrypt_AVX2)
    return 0;
}
